#include <pybind11/numpy.h>
#include <pybind11/pybind11.h>
#include <pybind11/stl.h>
#include <vector>
#include<iostream>
#include "mymath.h"
#include <math.h>
// #include <pybind11/eigen.h>
namespace py = pybind11;


std::vector<std::vector<std::vector<float>>> fft2d(int opt,int channel, std::vector<std::vector<std::vector<float>>>& src) {

    std::vector<std::vector<std::vector<float>>> out;

    // assert(a.size() == b.size());
    int ilen = src[0].size(),jlen = src[0][0].size();
    if (src.size() == channel){
        //空域格式
        
        std::vector<std::vector<float>> tmpMap;
        std::vector<float> tmpLine;

        for(int k = 0 ; k < channel ; k ++){
            tmpMap.clear();
            for(int i = 0 ; i < ilen ; i ++){
                tmpLine.clear();
                for(int j = 0 ; j < jlen ; j ++){
                    tmpLine.push_back(0);
                }
                tmpMap.push_back(tmpLine);
            }
            src.push_back(tmpMap);
        }

        printf("channels:%d hight:%d width:%d\n",src.size(),src[0].size(),src[1].size());
        printf("FFT2d pre_deal_Done\n");
        FFT2D(ilen,jlen,opt,src,out);
        printf("FFT2d Done\n");


    }else if ( src.size() == (channel<<1)){
        //频域ifft

        FFT2D(ilen,jlen,opt,src,out);

    }

    return out;

}

PYBIND11_MODULE(mypymath,m)
{
  m.doc() = "pybind11 example plugin";
  m.def("fft2d", &fft2d, py::return_value_policy::reference_internal );
  
}